library(tidyverse)
library(broom)
library(haven)
library(ggplot2)
library(ggrepel)
library(readr)
library(lme4)
library(interactions)
library(patchwork)
library(ggridges)
library(jtools)
library(modelsummary)
library(colorspace)
library(interplot)
library(ggpubr)
library(margins)
library(estimatr)
library(sandwich)
library(patchwork)
library(stargazer)

data2 <- read_dta("Study_2.dta")

# Create treatment variable
data2 <- data2 %>%
  mutate(treatment = case_when(
    trustexperimentsplit == 1 ~ 0,
    trustexperimentsplit == 2 ~ 1
  )) %>%
  mutate(treatment = factor(treatment, labels = c("Control", "Treatment")))

# Create trust variables
data2 <- data2 %>%
  mutate(
    trustMP = coalesce(trustExperimentGroup1MP, trustExperimentGroup2MP),
    trustMSP = coalesce(trustExperimentGroup1MSP, trustExperimentGroup2MSP),
    trustScotMin = coalesce(trustExperimentGroup1ScotMin, trustExperimentGroup2ScotMin),
    trustUKMin = coalesce(trustExperimentGroup1UKMin, trustExperimentGroup2UKMin),
    trustCivWest = coalesce(trustExperimentGroup1CivWest, trustExperimentGroup2CivWest),
    trustCivScot = coalesce(trustExperimentGroup1CivScot, trustExperimentGroup2CivScot)
  ) %>%
  mutate(across(c(trustMP, trustMSP, trustScotMin, trustUKMin, trustCivWest, trustCivScot),
                ~if_else(. == 8, NA_real_, . - 1)))

# Create binary incumbent_snp variable
data2 <- data2 %>% 
  mutate(incumbent_snp = case_when(
    voteIntentionWestminster == 4 ~ 1,
    TRUE ~ 0
  ))

# Create subsets for SNP voters and non-SNP voters
snp_voters <- subset(data2, incumbent_snp == 1)
non_snp_voters <- subset(data2, incumbent_snp == 0)

# Function to run models and create plots for each outcome
create_outcome_plot <- function(outcome_var, data_snp, data_non_snp) {
  # Determine overall range for y-axis
  y_min <- min(c(data_snp[[outcome_var]], data_non_snp[[outcome_var]]), na.rm = TRUE)
  y_max <- max(c(data_snp[[outcome_var]], data_non_snp[[outcome_var]]), na.rm = TRUE)
  
  # Calculate the range and add a small buffer (e.g., 5% of the range)
  y_range <- y_max - y_min
  buffer <- 0.10 * y_range
  
  y_limits <- c(max(0, y_min - buffer), min(5, y_max + buffer))
  
  # Models
  model_snp <- lm(as.formula(paste(outcome_var, "~ treatment")), data = data_snp)
  model_non_snp <- lm(as.formula(paste(outcome_var, "~ treatment")), data = data_non_snp)
  
  # Predictions
  data_snp$predicted <- predict(model_snp, data_snp)
  data_non_snp$predicted <- predict(model_non_snp, data_non_snp)
  
  # Subsets for treatment and control
  treat_snp <- subset(data_snp, treatment == "Treatment")
  control_snp <- subset(data_snp, treatment == "Control")
  treat_non_snp <- subset(data_non_snp, treatment == "Treatment")
  control_non_snp <- subset(data_non_snp, treatment == "Control")
  
  # Calculate ATE and p-values
  ate_snp <- coef(model_snp)["treatmentTreatment"]
  ate_non_snp <- coef(model_non_snp)["treatmentTreatment"]
  p_value_snp <- summary(model_snp)$coefficients["treatmentTreatment", "Pr(>|t|)"]
  p_value_non_snp <- summary(model_non_snp)$coefficients["treatmentTreatment", "Pr(>|t|)"]
  
  # Function to add asterisks based on p-value
  add_asterisks <- function(p_value) {
    if (p_value < 0.01) return("***")
    else if (p_value < 0.05) return("**")
    else if (p_value < 0.1) return("*")
    else return("")
  }
  
  # Create plots
  plot_snp <- effect_plot(model = model_snp, pred = treatment, robust = TRUE,
                          cat.geom = "point", cat.interval.geom = "linerange",
                          colors = "black", cat.pred.point.size = 3, int.width = .90) +
    labs(title = paste("SNP Voters (N =", nrow(data_snp), ")")) +
    ylab(paste("Trust in", gsub("trust", "", outcome_var))) +
    xlab("") +
    scale_x_discrete(labels = c("Control" = "Control", "Treatment" = "Treatment")) +
    scale_y_continuous(limits = y_limits) +
    geom_jitter(data = treat_snp, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, size = 3, width = .35, alpha = .15, shape = 20,
                pch = 21, color = "#11CE9E") +
    geom_jitter(data = control_snp, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 17,
                pch = 21, size = 3, color = "#CE1141FF") +
    geom_bracket(xmin = "Control", xmax = "Treatment",
                 y.position = y_limits[2] * 0.95,
                 label = sprintf("ATE=%.2f%s", ate_snp, add_asterisks(p_value_snp)),
                 tip.length = 0.02,
                 color = "black") +
    theme_minimal(base_size = 14) +
    theme(axis.title.y = element_text(face = "bold"),
          axis.text.x = element_text(face = "bold"))
  
  plot_non_snp <- effect_plot(model = model_non_snp, pred = treatment, robust = TRUE,
                              cat.geom = "point", cat.interval.geom = "linerange",
                              colors = "black", cat.pred.point.size = 3, int.width = .90) +
    labs(title = paste("Non-SNP Voters (N =", nrow(data_non_snp), ")")) +
    ylab("") +
    xlab("") +
    scale_x_discrete(labels = c("Control" = "Control", "Treatment" = "Treatment")) +
    scale_y_continuous(limits = y_limits) +
    geom_jitter(data = treat_non_snp, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 20,
                pch = 21, size = 3, color = "#11CE9E") +
    geom_jitter(data = control_non_snp, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 17,
                pch = 21, size = 3, color = "#CE1141FF") +
    geom_bracket(xmin = "Control", xmax = "Treatment",
                 y.position = y_limits[2] * 0.95,
                 label = sprintf("ATE=%.2f%s", ate_non_snp, add_asterisks(p_value_non_snp)),
                 tip.length = 0.02,
                 color = "black") +
    theme_minimal(base_size = 14) +
    theme(axis.title.y = element_blank(),
          axis.text.y = element_blank(),
          axis.text.x = element_text(face = "bold"))
  
  # Combine plots and return
  plot_snp + plot_non_snp +
    plot_annotation(title = paste('Effect of treatment on trust in', gsub("trust", "", outcome_var)),
                    theme = theme(plot.title = element_text(size = 14, face = "bold")))
}

# Create plots for each outcome
outcomes <- c("trustMP", "trustMSP", "trustUKMin", "trustScotMin", "trustCivWest", "trustCivScot")
plots <- map(outcomes, ~create_outcome_plot(.x, snp_voters, non_snp_voters))

# Combine all plots into one
combined_plot <- wrap_plots(plots, ncol = 2) +
  plot_annotation(
    title = 'Effects of treatment on trust in various entities',
    caption = "Treatment group outcome statistically distinct at p<.1(*), p<0.05(**), & p<0.01(***)",
    theme = theme(plot.title = element_text(size = 20, face = "bold"))
  )

# Save the combined plot
ggsave("Figure_9.png", combined_plot, width = 20, height = 30, dpi = 300)

# Create models
models <- list()
for (outcome in outcomes) {
  models[[outcome]] <- list(
    all = lm(as.formula(paste(outcome, "~ treatment")), data = data2),
    interaction = lm(as.formula(paste(outcome, "~ treatment * incumbent_snp")), data = data2),
    snp = lm(as.formula(paste(outcome, "~ treatment")), data = snp_voters),
    non_snp = lm(as.formula(paste(outcome, "~ treatment")), data = non_snp_voters)
  )
}

# Create and save tables
for (outcome in outcomes) {
  stargazer(models[[outcome]],
            type = "latex",
            title = paste("Models for", outcome),
            column.labels = c("All", "Interaction", "SNP", "Non-SNP"),
            out = paste0("Table_", outcome, "_incumbent_snp.tex"))
}